# 版本: 1.0.0
# 作者: 陈振玺
# 功能: 根据Excel文件中的FlowCell和index信息查找对应的fastq文件路径并计算文件大小

# 所有包导入必须放在文件头部
import argparse
import os
import glob
import pandas as pd
from pathlib import Path


class FastqFinder:
    """
    Fastq文件查找器
    根据FlowCell和index信息查找对应的fastq文件路径
    """
    
    def __init__(self, excel_file: str, fastq_dir: str):
        """
        初始化FastqFinder
        
        Args:
            excel_file: Excel文件路径，包含FlowCell和index列
            fastq_dir: fastq文件目录路径
        """
        self.excel_file = excel_file
        self.fastq_dir = fastq_dir
        self.data = None
    
    def load_excel(self):
        """
        加载Excel文件
        """
        try:
            self.data = pd.read_excel(self.excel_file)
            print(f"成功加载Excel文件: {self.excel_file}")
            print(f"数据行数: {len(self.data)}")
            
            # 检查必需的列
            required_columns = ['FlowCell', 'index']
            missing_columns = [col for col in required_columns if col not in self.data.columns]
            if missing_columns:
                raise ValueError(f"Excel文件缺少必需的列: {missing_columns}")
                
        except Exception as e:
            raise Exception(f"加载Excel文件失败: {str(e)}")
    
    def find_fastq_files(self, flowcell: str, index: str) -> list:
        """
        根据FlowCell和index查找对应的fastq文件
        
        Args:
            flowcell: FlowCell名称
            index: index名称
            
        Returns:
            list: 找到的fastq文件路径列表
        """
        found_files = []
        
        # 模式1: /mnt/data/{FlowCell}*/{FlowCell}_L*_{index}.fq.gz
        pattern1 = os.path.join(self.fastq_dir, f"{flowcell}*", f"{flowcell}_L*_{index}.fq.gz")
        files1 = glob.glob(pattern1)
        
        if files1:
            found_files.extend(files1)
        else:
            # 模式2: /mnt/data/{FlowCell}*/*/{FlowCell}_L*_{index}.fq.gz
            pattern2 = os.path.join(self.fastq_dir, f"{flowcell}*", "*", f"{flowcell}_L*_{index}.fq.gz")
            files2 = glob.glob(pattern2)
            
            if files2:
                found_files.extend(files2)
            else:
                # 模式3: /mnt/data/{FlowCell}*/*/{FlowCell}_L*_{index}_1.fq.gz
                pattern3 = os.path.join(self.fastq_dir, f"{flowcell}*", "*", f"{flowcell}_L*_{index}_1.fq.gz")
                files3 = glob.glob(pattern3)
                
                if files3:
                    found_files.extend(files3)
        
        return found_files
    
    def get_file_size_gb(self, file_path: str) -> float:
        """
        获取文件大小（GB）
        
        Args:
            file_path: 文件路径
            
        Returns:
            float: 文件大小（GB）
        """
        try:
            size_bytes = os.path.getsize(file_path)
            size_gb = size_bytes / (1024 ** 3)  # 转换为GB
            return round(size_gb, 3)
        except Exception:
            return 0.0
    
    def process_data(self):
        """
        处理数据，添加FqPath和Size列
        """
        if self.data is None:
            raise Exception("请先加载Excel文件")
        
        fq_paths = []
        sizes = []
        found_count = 0  # 统计找到文件的数量
        total_files = 0  # 统计总文件数
        
        for _, row in self.data.iterrows():
            flowcell = str(row['FlowCell'])
            index = str(row['index'])
            
            # 查找fastq文件
            found_files = self.find_fastq_files(flowcell, index)
            
            if not found_files:
                # 找不到文件
                fq_paths.append('')
                sizes.append('')
            elif len(found_files) == 1:
                # 找到一个文件
                file_path = found_files[0]
                fq_paths.append(file_path)
                sizes.append(self.get_file_size_gb(file_path))
                found_count += 1
                total_files += 1
            else:
                # 找到多个文件，每个文件一行
                file_paths = '\n'.join(found_files)
                fq_paths.append(file_paths)
                
                # 计算总大小
                total_size = sum(self.get_file_size_gb(f) for f in found_files)
                sizes.append(total_size)
                found_count += 1
                total_files += len(found_files)
        
        # 添加新列
        self.data['FqPath'] = fq_paths
        self.data['Size'] = sizes
        
        # 打印统计信息
        self.print_statistics(found_count, total_files, sizes)
    
    def print_statistics(self, found_count: int, total_files: int, sizes: list):
        """
        打印统计信息
        
        Args:
            found_count: 找到文件的样本数量
            total_files: 总文件数量
            sizes: 文件大小列表
        """
        import numpy as np
        
        print("\n" + "=" * 50)
        print("统计信息")
        print("=" * 50)
        
        # 路径统计
        total_samples = len(self.data)
        not_found_count = total_samples - found_count
        
        print(f"总样本数: {total_samples}")
        print(f"找到文件的样本数: {found_count} ({found_count/total_samples*100:.1f}%)")
        print(f"未找到文件的样本数: {not_found_count} ({not_found_count/total_samples*100:.1f}%)")
        print(f"总文件数: {total_files}")
        
        # Size统计（只统计有效的数值）
        valid_sizes = [s for s in sizes if isinstance(s, (int, float)) and s > 0]
        
        if valid_sizes:
            print("\nSize统计 (GB):")
            print(f"  样本数: {len(valid_sizes)}")
            print(f"  最小值: {min(valid_sizes):.3f}")
            print(f"  最大值: {max(valid_sizes):.3f}")
            print(f"  平均值: {np.mean(valid_sizes):.3f}")
            print(f"  中位数: {np.median(valid_sizes):.3f}")
            print(f"  标准差: {np.std(valid_sizes):.3f}")
            print(f"  总大小: {sum(valid_sizes):.3f}")
            
            # 四分位数
            q25 = np.percentile(valid_sizes, 25)
            q75 = np.percentile(valid_sizes, 75)
            print(f"  25%分位数: {q25:.3f}")
            print(f"  75%分位数: {q75:.3f}")
        else:
            print("\n没有有效的Size数据")
        
        print("=" * 50)
    
    def save_result(self, output_file: str):
        """
        保存结果到文件
        
        Args:
            output_file: 输出文件路径
        """
        if self.data is None:
            raise Exception("没有数据可保存")
        
        try:
            # 根据文件扩展名选择保存格式
            file_ext = Path(output_file).suffix.lower()
            
            if file_ext == '.xlsx':
                self.data.to_excel(output_file, index=False)
            elif file_ext == '.csv':
                self.data.to_csv(output_file, index=False)
            else:
                # 默认保存为Excel格式
                self.data.to_excel(output_file, index=False)
            
            print(f"结果已保存到: {output_file}")
            print(f"处理完成，共处理 {len(self.data)} 行数据")
            
        except Exception as e:
            raise Exception(f"保存文件失败: {str(e)}")


def main():
    """
    主函数
    """
    parser = argparse.ArgumentParser(
        description='根据Excel文件中的FlowCell和index信息查找对应的fastq文件路径',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  fastq-finder -i input.xlsx -f /mnt/data/ -o output.xlsx
  fastq-finder -i sample.xlsx -f /path/to/fastq/ -o result.csv

输入文件格式:
  Excel文件必须包含以下列:
  - FlowCell: 流动池名称
  - index: 索引序列

输出文件格式:
  在原有列基础上添加:
  - FqPath: 找到的fastq文件路径（多个文件时每行一个）
  - Size: 文件大小（GB）
        """
    )
    
    parser.add_argument('-i', '--input', required=True,
                       help='输入Excel文件路径，包含FlowCell和index列')
    parser.add_argument('-f', '--fastq-dir', required=True,
                       help='fastq文件目录路径，如 /mnt/data/')
    parser.add_argument('-o', '--output', required=True,
                       help='输出文件路径')
    
    args = parser.parse_args()
    
    try:
        # 检查输入文件是否存在
        if not os.path.exists(args.input):
            raise FileNotFoundError(f"输入文件不存在: {args.input}")
        
        # 检查fastq目录是否存在
        if not os.path.exists(args.fastq_dir):
            raise FileNotFoundError(f"fastq目录不存在: {args.fastq_dir}")
        
        # 创建输出目录（如果不存在）
        output_dir = os.path.dirname(args.output)
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        print(f"开始处理...")
        print(f"输入文件: {args.input}")
        print(f"fastq目录: {args.fastq_dir}")
        print(f"输出文件: {args.output}")
        print("-" * 50)
        
        # 创建FastqFinder实例并处理
        finder = FastqFinder(args.input, args.fastq_dir)
        finder.load_excel()
        finder.process_data()
        finder.save_result(args.output)
        
        print("-" * 50)
        print("处理完成！")
        
    except Exception as e:
        print(f"错误: {str(e)}")
        return 1
    
    return 0


if __name__ == '__main__':
    exit(main())